# from model.semseg.segformer import Segformer
# import torch
# import yaml
# with open('configs/pascal_segformer.yaml') as f:
#     cfg = yaml.load(f, yaml.Loader)
# model = Segformer(cfg)
# print(model)
# x = torch.rand(2, 3, 224, 224)
# out = model(x)
# print(out.shape)


# import fastvit
# import fastvit
# from fastvit.models import fastvit_ma36
# from timm import create_model

# # model = fastvit_ma36()
# model = create_model('fastvit_t8')
# print(model)

import fastvit
from fastvit import models
from fastvit.models.modules.mobileone import reparameterize_model
import timm
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torchvision.datasets import ImageFolder, ImageNet
from torchvision import transforms as TTF
from tqdm import tqdm
from tlhengine.utils import Timer
import argparse

def parse_args():
    parser = argparse.ArgumentParser('imagenet validation')
    
def evaluate(model, val_loader, device):
    model.eval()
    total = correct = 0
    with torch.no_grad():
        for images, labels in tqdm(val_loader, ncols=10):
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    accuracy = 100 * correct / total
    print(f'Validation accuracy: {accuracy:.2f}%')
    
device = 'cuda:1'
model : nn.Module= timm.create_model('fastvit_ma36').to(device)
# # model = timm.create_model('fast')

ckpt = torch.load('/root/code/apple/fastvit/ckpt/fastvit_ma36.pth.tar')
# # print(ckpt)
with Timer('loading weights'):
    model.load_state_dict(ckpt['state_dict'])
# print('##########################')
# print(model)

tranf = TTF.Compose([
    TTF.Resize(256),
    TTF.CenterCrop(224),
    TTF.ToTensor(),
    TTF.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
val_set = ImageNet('/root/data/imagenet', split='val', transform=tranf)
val_loader = DataLoader(val_set, batch_size=32, num_workers=8, persistent_workers=True, pin_memory=True, pin_memory_device=device)
total = correct = 0

with Timer('eval before repar'):
    evaluate(model, val_loader, device)

with Timer('reparamerization'):
    model_inf = reparameterize_model(model) 
with Timer('eval after repar'):
    evaluate(model_inf, val_loader, device)

    
print('done')